// File: src/main/java/com/example/chatweb/service/impl/ChatServiceImpl.java

package com.example.chatweb.service.impl;
import org.json.JSONArray;
import org.json.JSONObject;

import com.example.chatweb.entity.Message;
import com.example.chatweb.service.ChatService;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import java.util.*;


@Service
public class ChatServiceImpl implements ChatService {

    private final RestTemplate restTemplate;

    @Autowired
    public ChatServiceImpl(RestTemplate restTemplate) {
        this.restTemplate = restTemplate;
    }

    @Override
    public String getChat(Message message) {
        System.out.println(message.toString());
        System.out.println(message.getModelName());
        if ("MedicalGPT".equals(message.getModelName())){
            String url = "http://10.112.203.120:8001/v1/chat/completions";

            JSONObject requestBody = new JSONObject();
            requestBody.put("model", message.getModelName());

            JSONArray messages = new JSONArray();
            JSONObject singleMessage = new JSONObject();
            singleMessage.put("content", message.getQuery());
            singleMessage.put("role", "user");
            messages.put(singleMessage);

            requestBody.put("messages", messages);
            requestBody.put("tools", new JSONArray());
            requestBody.put("do_sample", true);
            requestBody.put("temperature", 0.95);
            requestBody.put("top_p", 0.7);
            requestBody.put("n", 1);
            requestBody.put("max_tokens", 512);
            requestBody.put("stream", true);

            // 将JSONObject转换为String
            String json = requestBody.toString();

            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);

            // 创建 HttpEntity 对象，此处使用String作为body
            HttpEntity<String> request = new HttpEntity<>(json, headers);

            // 发送POST请求
            String response = restTemplate.postForObject(url, request, String.class);
            System.out.println(response);

            // 处理响应
            ObjectMapper objectMapper = new ObjectMapper();
            try {
                JsonNode rootNode = objectMapper.readTree(response);
                JsonNode textNode = rootNode.path("text");
                JsonNode choiceNode = rootNode.path("choices");
                if (!textNode.isMissingNode()) { // 如果"text"节点存在
                    return textNode.asText();
                } else if (!choiceNode.isEmpty() && choiceNode.isArray()) {
                    JsonNode firstChoice = choiceNode.get(0);
                    JsonNode res = firstChoice.path("message");
                    return res.path("content").asText();
                }
            } catch (Exception e) {
                return null;
            }
        }
        else if("BenTsao".equals(message.getModelName())){
            String url = "http://10.112.203.120:8007/v1/chat/completions";

            JSONObject requestBody = new JSONObject();
            requestBody.put("model", message.getModelName());

            JSONArray messages = new JSONArray();
            JSONObject singleMessage = new JSONObject();
            singleMessage.put("content", message.getQuery());
            singleMessage.put("role", "user");
            messages.put(singleMessage);

            requestBody.put("messages", messages);
            requestBody.put("tools", new JSONArray());
            requestBody.put("do_sample", true);
            requestBody.put("temperature", 0.95);
            requestBody.put("top_p", 0.7);
            requestBody.put("n", 1);
            requestBody.put("max_tokens", 512);
            requestBody.put("stream", false);

            // 将JSONObject转换为String
            String json = requestBody.toString();

            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);

            // 创建 HttpEntity 对象，此处使用String作为body
            HttpEntity<String> request = new HttpEntity<>(json, headers);

            // 发送POST请求
            String response = restTemplate.postForObject(url, request, String.class);
            System.out.println(response);

            // 处理响应
            ObjectMapper objectMapper = new ObjectMapper();
            try {
                JsonNode rootNode = objectMapper.readTree(response);
                JsonNode textNode = rootNode.path("text");
                JsonNode choiceNode = rootNode.path("choices");
                if (!textNode.isMissingNode()) { // 如果"text"节点存在
                    return textNode.asText();
                } else if (!choiceNode.isEmpty() && choiceNode.isArray()) {
                    JsonNode firstChoice = choiceNode.get(0);
                    JsonNode res = firstChoice.path("message");
                    return res.path("content").asText();
                }
            } catch (Exception e) {
                return null;
            }
        }
        else{
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            HttpEntity<Message> request = new HttpEntity<>(message, headers);
            String url = "http://10.112.203.120:7861/chat/chat";


            // 创建 HttpEntity 对象，包含消息对象和头部信息

            // 发送 POST 请求到目标接口，并接收返回的响应
            // 假设响应体是字符串类型
            String response = restTemplate.postForObject(url, request, String.class);
            String jsonPart = response.substring(6);
            ObjectMapper objectMapper = new ObjectMapper();
            try {
                JsonNode rootNode = objectMapper.readTree(jsonPart);
                String text = rootNode.get("text").asText();
                return text;
            } catch (Exception e) {
                e.printStackTrace();
                return null; // 解析失败时返回 null 或适当处理错误
            }
        }

        return null;
    }

    @Override
    public String getKnowledgeBaseChat(String message){
        String url = "http://10.112.203.120:7861/chat/knowledge_base_chat";
        JSONObject requestBody = new JSONObject();
        JSONObject query_t = new JSONObject(message.toString());
        String query = query_t.getString("message");
        requestBody.put("query", query);
        requestBody.put("knowledge_base_name", "X态势");
        requestBody.put("top_k", 3);
        requestBody.put("score_threshold", 1);
        JSONArray messages = new JSONArray();
//        JSONObject singleMessage = new JSONObject();
//        singleMessage.put("content", "");
//        singleMessage.put("role", "");
//        messages.put(singleMessage);
        requestBody.put("stream", false);
        requestBody.put("history", messages);
        requestBody.put("model_name", "chatglm3-6b");
        requestBody.put("temperature", 0.70);
        requestBody.put("max_tokens", 0);
        requestBody.put("prompt_name", "default");
        String json = requestBody.toString();

        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);

        // 创建 HttpEntity 对象，此处使用String作为body
        HttpEntity<String> request = new HttpEntity<>(json, headers);

        // 发送POST请求
        String response = restTemplate.postForObject(url, request, String.class);
//        System.out.println(response);
        System.out.println(requestBody);
        System.out.println(response);
        String jsonPart = response.startsWith("data: ") ? response.substring(6): response;
        // 处理响应
        ObjectMapper objectMapper = new ObjectMapper();
        try {
            System.out.println(response);
            JsonNode rootNode = objectMapper.readTree(jsonPart);
            JsonNode textNode = rootNode.path("answer");
            if(!textNode.path("message").isMissingNode()){
                textNode = textNode.path("message");
            }
            System.out.println(textNode.asText());
            JsonNode docsNode = rootNode.path("docs");
            List<String> documents = new ArrayList<>();
            StringBuilder sources = new StringBuilder();
            if (!docsNode.isMissingNode() && docsNode.isArray()) {
                for (JsonNode docNode : docsNode) {
                    String docText = docNode.asText();
                    // Split the document text by new lines
                    String[] lines = docText.split("\n");
                    for (String line : lines) {
                        if (line.startsWith("出处")) {
                            sources.append(line.trim()).append("\n");
                        }
                    }
                }
            }
            System.out.println("1111");
            System.out.println(docsNode);
            System.out.println("1111");
            System.out.println(sources);
            System.out.println("22222");
            return textNode.asText()+'\n'+sources;
        } catch (Exception e) {
            return null;
        }

    }
    @Override
    public String Kchat(Message message){
        return  null;
    }
    }
